#  -*- coding: utf-8 -*-

from pymongo import UpdateOne,ASCENDING, DESCENDING
from factor.base_factor import BaseFactor
from data.finance_report_crawler import FinanceReportCrawler
from data.data_module import DataModule
from util.stock_util import get_all_codes,get_all_indexes_date,calc_negative_diff_dates,multi_computer,get_trading_dates,get_all_codes_trading_date,get_all_indexes
from util.database import DB_CONN
import time
import pandas as pd
import numpy as np
from datetime import datetime, timedelta


"""
实现后复权的相对强度（RS）一类因子的计算
"""


class RSFactor(BaseFactor):
    def __init__(self):
        BaseFactor.__init__(self, name='rs')

    def compute_single(self,code,is_index, autype, begin_date,end_date):
        start_time = time.time()
        dm = DataModule()
        print('计算相对涨幅, %s' % code)

        if begin_date is None:
            begin_date = "1990-12-19"

        df_daily = dm.get_k_data(code, index=is_index, autype=autype,begin_date = None,end_date= end_date)
        if df_daily.index.size > 0:
            df_daily.set_index(['date'], 1, inplace=True)
            update_requests = []
            df_trading_daily = df_daily.loc[df_daily.is_trading == True, :]
            df_trading_daily_copy = df_trading_daily.copy()
            try:
                # 计算相对涨幅
                df_trading_daily_copy['rs_change_rate_5'] = df_trading_daily_copy['close'].rolling(5).apply(lambda x: round(100*(x[4] - x[0])/x[0],2))
                df_trading_daily_copy['rs_change_rate_10'] = df_trading_daily_copy['close'].rolling(10).apply(lambda x: round(100*(x[9] - x[0])/x[0],2))
                df_trading_daily_copy['rs_change_rate_20'] = df_trading_daily_copy['close'].rolling(20).apply(lambda x: round(100*(x[19] - x[0])/x[0],2))
                df_trading_daily_copy['rs_change_rate_60'] = df_trading_daily_copy['close'].rolling(60).apply(lambda x: round(100*(x[59] - x[0])/x[0],2))
                df_trading_daily_copy['rs_change_rate_120'] = df_trading_daily_copy['close'].rolling(120).apply(lambda x: round(100*(x[119] - x[0])/x[0],2))
                df_trading_daily_copy['rs_change_rate_250'] = df_trading_daily_copy['close'].rolling(250).apply(lambda x: round(100*(x[249] - x[0])/x[0],2))

                df_trading_daily_copy['rs_mean_amount_5'] = round(df_trading_daily_copy['amount'].rolling(5).mean()/100000000,2)
                df_trading_daily_copy['rs_mean_amount_10'] = round(df_trading_daily_copy['amount'].rolling(10).mean()/100000000,2)
                df_trading_daily_copy['rs_mean_amount_20'] = round(df_trading_daily_copy['amount'].rolling(20).mean()/100000000,2)
                df_trading_daily_copy['rs_mean_amount_60'] = round(df_trading_daily_copy['amount'].rolling(60).mean()/100000000,2)
                df_trading_daily_copy['rs_mean_amount_120'] = round(df_trading_daily_copy['amount'].rolling(120).mean()/100000000,2)
                df_trading_daily_copy['rs_mean_amount_250'] = round(df_trading_daily_copy['amount'].rolling(250).mean()/100000000,2)
                #print(df_trading_daily_copy)

                df_target_daily = df_trading_daily_copy.loc[begin_date:end_date]
                #print(df_target_daily)
                for date in df_target_daily.index:
                    update_requests.append(
                        UpdateOne(
                            {'code': code, 'date': date,'index': is_index},
                            {'$set': {'code': code,
                                      'date': date,
                                      'index': is_index,
                                      'rs_change_rate_5': df_target_daily.loc[date]['rs_change_rate_5'],
                                      'rs_change_rate_10': df_target_daily.loc[date]['rs_change_rate_10'],
                                      'rs_change_rate_20': df_target_daily.loc[date]['rs_change_rate_20'],
                                      'rs_change_rate_60': df_target_daily.loc[date]['rs_change_rate_60'],
                                      'rs_change_rate_120': df_target_daily.loc[date]['rs_change_rate_120'],
                                      'rs_change_rate_250': df_target_daily.loc[date]['rs_change_rate_250'],
                                      'rs_mean_amount_5': df_target_daily.loc[date]['rs_mean_amount_5'],
                                      'rs_mean_amount_10': df_target_daily.loc[date]['rs_mean_amount_10'],
                                      'rs_mean_amount_20': df_target_daily.loc[date]['rs_mean_amount_20'],
                                      'rs_mean_amount_60': df_target_daily.loc[date]['rs_mean_amount_60'],
                                      'rs_mean_amount_120': df_target_daily.loc[date]['rs_mean_amount_120'],
                                      'rs_mean_amount_250': df_target_daily.loc[date]['rs_mean_amount_250'],
                                      }},
                            upsert=True))
            except:
                print('填充相对涨幅时发生错误，股票代码：%s 是否指数：%s' % (code,is_index), flush=True)
            #
            if len(update_requests) > 0:
                update_result = self.collection.bulk_write(update_requests, ordered=False)

                end_time = time.time()
                print('填充相对涨幅，股票：%s，是否指数：%s,插入：%4d条，更新：%4d条,耗时：%.3f 秒' %
                      (code, is_index, update_result.upserted_count, update_result.modified_count, (end_time - start_time)),
                      flush=True)

    def compute_rs_date(self,date,is_index):
        start_time = time.time()
        update_requests = []

        codes = get_all_codes_trading_date(date)
        codes_count = len(codes)
        rs_date_df = pd.DataFrame(columns=["rs_change_rate_5","rs_change_rate_10","rs_change_rate_20","rs_5","rs_10","rs_20",
            "rs_change_rate_60","rs_change_rate_120","rs_change_rate_250","rs_60","rs_120","rs_250"],index=codes)
        i = 0
        for code in codes:
            factor_cursor = DB_CONN['rs'].find(
                {'code': code, 'date':date, "index": is_index},
                projection={'_id': False})
            for factor in factor_cursor:
                #判断存在250日相对涨幅的才进行赋值，这样为了可以过滤排除上市1年以内的新股
                if pd.notnull(factor["rs_change_rate_250"]):
                    rs_date_df.loc[code]["rs_change_rate_5"] = factor['rs_change_rate_5']
                    rs_date_df.loc[code]["rs_change_rate_10"] = factor['rs_change_rate_10']
                    rs_date_df.loc[code]["rs_change_rate_20"] = factor['rs_change_rate_20']
                    rs_date_df.loc[code]["rs_change_rate_60"] = factor['rs_change_rate_60']
                    rs_date_df.loc[code]["rs_change_rate_120"] = factor['rs_change_rate_120']
                    rs_date_df.loc[code]["rs_change_rate_250"] = factor['rs_change_rate_250']

        #print(rs_date_df)
        #丢弃一行数据全部都是NaN的行，即1年以内的新股数据
        rs_date_df.dropna(how='all',inplace=True)

        # 从小到大计算排名
        rs_date_df['rs_5_rank'] = rs_date_df["rs_change_rate_5"].rank()
        rs_date_df['rs_10_rank'] = rs_date_df["rs_change_rate_10"].rank()
        rs_date_df['rs_20_rank'] = rs_date_df["rs_change_rate_20"].rank()
        rs_date_df['rs_60_rank'] = rs_date_df["rs_change_rate_60"].rank()
        rs_date_df['rs_120_rank'] = rs_date_df["rs_change_rate_120"].rank()
        rs_date_df['rs_250_rank'] = rs_date_df["rs_change_rate_250"].rank()

        #按百分位归类
        for i in range(100):
            ratio_min = rs_date_df.index.size * i / 100
            ratio_max = rs_date_df.index.size* (i+1) / 100
            rs_date_df.loc[(rs_date_df['rs_5_rank'] <= ratio_max) & (rs_date_df['rs_5_rank'] > ratio_min),'rs_5'] = i
            rs_date_df.loc[(rs_date_df['rs_10_rank'] <= ratio_max) & (rs_date_df['rs_10_rank'] > ratio_min),'rs_10'] = i
            rs_date_df.loc[(rs_date_df['rs_20_rank'] <= ratio_max) & (rs_date_df['rs_20_rank'] > ratio_min),'rs_20'] = i
            rs_date_df.loc[(rs_date_df['rs_60_rank'] <= ratio_max) & (rs_date_df['rs_60_rank'] > ratio_min),'rs_60'] = i
            rs_date_df.loc[(rs_date_df['rs_120_rank'] <= ratio_max) & (rs_date_df['rs_120_rank'] > ratio_min),'rs_120'] = i
            rs_date_df.loc[(rs_date_df['rs_250_rank'] <= ratio_max) & (rs_date_df['rs_250_rank'] > ratio_min),'rs_250'] = i

        rs_date_df.fillna(-1,inplace=True)

        for code in rs_date_df.index:
            update_requests.append(
                UpdateOne(
                    {'code': code, 'date': date, 'index': is_index},
                    {'$set': {
                              'rs_5': rs_date_df.loc[code]['rs_5'],
                              'rs_10': rs_date_df.loc[code]['rs_10'],
                              'rs_20': rs_date_df.loc[code]['rs_20'],
                              'rs_60': rs_date_df.loc[code]['rs_60'],
                              'rs_120': rs_date_df.loc[code]['rs_120'],
                              'rs_250': rs_date_df.loc[code]['rs_250']}},
                    upsert=True))


        if len(update_requests) > 0:
            update_result = self.collection.bulk_write(update_requests, ordered=False)
            end_time = time.time()
            print('填充RS因子，日期：%s，插入：%4d条，更新：%4d条,耗时：%.3f 秒' %
                  (date, update_result.upserted_count, update_result.modified_count, (end_time - start_time)),
                  flush=True)

    def compute_single_index(self,code,is_index, autype, begin_date,end_date):
        start_time = time.time()
        dm = DataModule()
        print('计算相对涨幅, %s' % code)

        if begin_date is None:
            begin_date = "1990-12-19"

        df_daily = dm.get_k_data(code, index=is_index, autype=autype,begin_date = None,end_date= end_date)
        if df_daily.index.size > 0:
            df_daily.set_index(['date'], 1, inplace=True)
            update_requests = []
            df_trading_daily = df_daily.loc[df_daily.is_trading == True, :]
            df_trading_daily_copy = df_trading_daily.copy()
            try:
                # 计算相对涨幅
                for i in [5,10,20,60,120,250]:
                    if df_daily.index.size >= i:
                        df_trading_daily_copy[f'rs_change_rate_{i}'] = df_trading_daily_copy['close'].rolling(i).apply(lambda x: round(100*(x[i-1] - x[0])/x[0],2),raw=True)
                    else:
                        df_trading_daily_copy[f'rs_change_rate_{i}'] = np.nan
                #print(df_trading_daily_copy)

                df_target_daily = df_trading_daily_copy.loc[begin_date:end_date]
                #print(df_target_daily)
                for date in df_target_daily.index:
                    update_requests.append(
                        UpdateOne(
                            {'code': code, 'date': date,'index': is_index},
                            {'$set': {'code': code,
                                      'date': date,
                                      'index': is_index,
                                      'rs_change_rate_5': df_target_daily.loc[date]['rs_change_rate_5'],
                                      'rs_change_rate_10': df_target_daily.loc[date]['rs_change_rate_10'],
                                      'rs_change_rate_20': df_target_daily.loc[date]['rs_change_rate_20'],
                                      'rs_change_rate_60': df_target_daily.loc[date]['rs_change_rate_60'],
                                      'rs_change_rate_120': df_target_daily.loc[date]['rs_change_rate_120'],
                                      'rs_change_rate_250': df_target_daily.loc[date]['rs_change_rate_250']}},
                            upsert=True))
            except:
                print('填充相对涨幅时发生错误，股票代码：%s 是否指数：%s' % (code,is_index), flush=True)
            #
            if len(update_requests) > 0:
                update_result = self.collection.bulk_write(update_requests, ordered=False)

                end_time = time.time()
                print('填充相对涨幅，股票：%s，是否指数：%s,插入：%4d条，更新：%4d条,耗时：%.3f 秒' %
                      (code, is_index, update_result.upserted_count, update_result.modified_count, (end_time - start_time)),
                      flush=True)

    def compute_rs_index_date(self,date,is_index):
        start_time = time.time()
        update_requests = []
        codes = []
        codes = get_all_indexes_date(date)
        codes_count = len(codes)
        rs_date_df = pd.DataFrame(columns=["rs_change_rate_5","rs_change_rate_10","rs_change_rate_20","rs_5","rs_10","rs_20",
            "rs_change_rate_60","rs_change_rate_120","rs_change_rate_250","rs_60","rs_120","rs_250"],index=codes)
        i = 0
        for code in codes:
            #print(code)
            factor_cursor = DB_CONN['rs'].find(
                {'code': code, 'date':date, "index": is_index},
                projection={'_id': False})
            for factor in factor_cursor:
                try:
                    rs_date_df.loc[code]["rs_change_rate_5"] = factor['rs_change_rate_5']
                    rs_date_df.loc[code]["rs_change_rate_10"] = factor['rs_change_rate_10']
                    rs_date_df.loc[code]["rs_change_rate_20"] = factor['rs_change_rate_20']
                    rs_date_df.loc[code]["rs_change_rate_60"] = factor['rs_change_rate_60']
                    rs_date_df.loc[code]["rs_change_rate_120"] = factor['rs_change_rate_120']
                    rs_date_df.loc[code]["rs_change_rate_250"] = factor['rs_change_rate_250']
                except:
                    rs_date_df.loc[code]["rs_change_rate_5"] = None
                    rs_date_df.loc[code]["rs_change_rate_10"] = None
                    rs_date_df.loc[code]["rs_change_rate_20"] = None
                    rs_date_df.loc[code]["rs_change_rate_60"] = None
                    rs_date_df.loc[code]["rs_change_rate_120"] = None
                    rs_date_df.loc[code]["rs_change_rate_250"] = None
        #print(rs_date_df)

        rs_date_df.dropna(how='all',inplace=True)

        # 从小到大计算排名
        rs_date_df['rs_5_rank'] = rs_date_df["rs_change_rate_5"].rank()
        rs_date_df['rs_10_rank'] = rs_date_df["rs_change_rate_10"].rank()
        rs_date_df['rs_20_rank'] = rs_date_df["rs_change_rate_20"].rank()
        rs_date_df['rs_60_rank'] = rs_date_df["rs_change_rate_60"].rank()
        rs_date_df['rs_120_rank'] = rs_date_df["rs_change_rate_120"].rank()
        rs_date_df['rs_250_rank'] = rs_date_df["rs_change_rate_250"].rank()

        #按百分位归类
        for i in range(100):
            ratio_min = rs_date_df.index.size * i / 100
            ratio_max = rs_date_df.index.size* (i+1) / 100
            #print(i,ratio_min,ratio_max)
            rs_date_df.loc[(rs_date_df['rs_5_rank'] <= ratio_max) & (rs_date_df['rs_5_rank'] > ratio_min),'rs_5'] = i
            rs_date_df.loc[(rs_date_df['rs_10_rank'] <= ratio_max) & (rs_date_df['rs_10_rank'] > ratio_min),'rs_10'] = i
            rs_date_df.loc[(rs_date_df['rs_20_rank'] <= ratio_max) & (rs_date_df['rs_20_rank'] > ratio_min),'rs_20'] = i
            rs_date_df.loc[(rs_date_df['rs_60_rank'] <= ratio_max) & (rs_date_df['rs_60_rank'] > ratio_min),'rs_60'] = i
            rs_date_df.loc[(rs_date_df['rs_120_rank'] <= ratio_max) & (rs_date_df['rs_120_rank'] > ratio_min),'rs_120'] = i
            rs_date_df.loc[(rs_date_df['rs_250_rank'] <= ratio_max) & (rs_date_df['rs_250_rank'] > ratio_min),'rs_250'] = i

        #rs_date_df.to_csv("index_rs_1113.csv")

        rs_date_df.fillna(-1,inplace=True)
        for code in rs_date_df.index:
            update_requests.append(
                UpdateOne(
                    {'code': code, 'date': date, 'index': is_index},
                    {'$set': {
                              'rs_5': rs_date_df.loc[code]['rs_5'],
                              'rs_10': rs_date_df.loc[code]['rs_10'],
                              'rs_20': rs_date_df.loc[code]['rs_20'],
                              'rs_60': rs_date_df.loc[code]['rs_60'],
                              'rs_120': rs_date_df.loc[code]['rs_120'],
                              'rs_250': rs_date_df.loc[code]['rs_250']}},
                    upsert=True))


        if len(update_requests) > 0:
            update_result = self.collection.bulk_write(update_requests, ordered=False)
            end_time = time.time()
            print('填充RS因子，日期：%s，插入：%4d条，更新：%4d条,耗时：%.3f 秒' %
                  (date, update_result.upserted_count, update_result.modified_count, (end_time - start_time)),
                  flush=True)


    def compute(self, begin_date, end_date):
        """
        计算指定时间段内所有股票的该因子的值(相对强度，并保存到数据库中
        60/120/250
        :param begin_date:  开始时间
        :param end_date: 结束时间
        """
        self.collection.create_index([('code', 1), ('date', 1),('index',1)])
        print("begin compute rs codes")

        #获取所有股票
        codes = get_all_codes()
        args = (False,"hfq",begin_date,end_date)
        multi_computer(compute_codes,codes,args)
        dates = get_trading_dates(begin_date,end_date)
        args = (False,)
        multi_computer(compute_dates,dates,args)

        #获取所有指数
        codes = get_all_indexes()
        args = (True,None,begin_date,end_date)
        multi_computer(compute_codes,codes,args)
        dates = get_trading_dates(begin_date,end_date)
        args = (True,)
        multi_computer(compute_dates,dates,args)


def compute_dates(dates,is_index):
    if is_index:
        for date in dates:
            RSFactor().compute_rs_index_date(date, is_index)
    else:
        for date in dates:
            RSFactor().compute_rs_date(date,is_index)

def compute_codes(codes,is_index,autype,begin_date,end_date):
    if is_index:
        for code in codes:
            RSFactor().compute_single_index(code,is_index,autype,begin_date,end_date)
    else:
        for code in codes:
            RSFactor().compute_single(code,is_index,autype,begin_date,end_date)


if __name__ == '__main__':
    # 执行因子的提取任务
    pd.set_option('display.max_columns', 1000)
    pd.set_option('display.width', 1000)
    rs = RSFactor()
    #rs.computer_single("002463",False,"hfq",begin_date="2016-01-01",end_date="2018-11-16")
    #rs.compute(begin_date="2018-01-01",end_date=None)
    rs.compute_single(code="600091",is_index=False,autype="hfq",begin_date="2018-01-01",end_date=None)
    #RSFactor().compute('1990-12-19', '2018-11-18')

    # dates = get_trading_dates(None, None)
    # for date in dates:
    #     #print(date)
    #     factor_cursor = DB_CONN['rs'].find(
    #         {'date': date, "index": False},
    #         projection={'_id': False})
    #     for factor in factor_cursor:
    #         if "rs_120" not in factor.keys():
    #             print(date,factor['code'])

